package com.github.codeql.utils

import com.github.codeql.KotlinUsesExtractor
import com.github.codeql.Logger
import com.github.codeql.getJavaEquivalentClassId
import com.github.codeql.utils.versions.codeQlWithHasQuestionMark
import com.github.codeql.utils.versions.createImplicitParameterDeclarationWithWrappedDescriptor
import org.jetbrains.kotlin.backend.common.extensions.IrPluginContext
import org.jetbrains.kotlin.descriptors.ClassKind
import org.jetbrains.kotlin.ir.builders.declarations.addConstructor
import org.jetbrains.kotlin.ir.builders.declarations.buildClass
import org.jetbrains.kotlin.ir.declarations.IrClass
import org.jetbrains.kotlin.ir.declarations.IrTypeParameter
import org.jetbrains.kotlin.ir.declarations.IrTypeParametersContainer
import org.jetbrains.kotlin.ir.declarations.impl.IrExternalPackageFragmentImpl
import org.jetbrains.kotlin.ir.declarations.impl.IrFactoryImpl
import org.jetbrains.kotlin.ir.expressions.IrConstructorCall
import org.jetbrains.kotlin.ir.expressions.impl.IrConstructorCallImpl
import org.jetbrains.kotlin.ir.symbols.IrTypeParameterSymbol
import org.jetbrains.kotlin.ir.symbols.impl.DescriptorlessExternalPackageFragmentSymbol
import org.jetbrains.kotlin.ir.types.*
import org.jetbrains.kotlin.ir.types.impl.IrSimpleTypeImpl
import org.jetbrains.kotlin.ir.types.impl.IrStarProjectionImpl
import org.jetbrains.kotlin.ir.types.impl.makeTypeProjection
import org.jetbrains.kotlin.ir.util.classId
import org.jetbrains.kotlin.ir.util.constructedClassType
import org.jetbrains.kotlin.ir.util.constructors
import org.jetbrains.kotlin.ir.util.kotlinFqName
import org.jetbrains.kotlin.ir.util.parents
import org.jetbrains.kotlin.name.FqName
import org.jetbrains.kotlin.name.Name
import org.jetbrains.kotlin.types.Variance
import org.jetbrains.kotlin.utils.addToStdlib.firstIsInstanceOrNull

fun IrType.substituteTypeArguments(params: List<IrTypeParameter>, arguments: List<IrTypeArgument>) =
    when (this) {
        is IrSimpleType -> substituteTypeArguments(params.map { it.symbol }.zip(arguments).toMap())
        else -> this
    }

private fun IrSimpleType.substituteTypeArguments(
    substitutionMap: Map<IrTypeParameterSymbol, IrTypeArgument>
): IrSimpleType {
    if (substitutionMap.isEmpty()) return this

    val newArguments =
        arguments.map {
            if (it is IrTypeProjection) {
                val itType = it.type
                if (itType is IrSimpleType) {
                    subProjectedType(substitutionMap, itType, it.variance)
                } else {
                    it
                }
            } else {
                it
            }
        }

    return IrSimpleTypeImpl(classifier, isNullable(), newArguments, annotations)
}

/**
 * Returns true if substituting `innerVariance T` into the context `outerVariance []` discards all
 * knowledge about what T could be.
 *
 * Note this throws away slightly more information than it could: for example, the projection "in
 * (out List)" can refer to any superclass of anything that implements List, which specifically
 * excludes e.g. String, but can't be represented as a type projection. The projection "out (in
 * List)" on the other hand really is equivalent to "out Any?", which is to say no bound at all.
 */
private fun conflictingVariance(outerVariance: Variance, innerVariance: Variance) =
    (outerVariance == Variance.IN_VARIANCE && innerVariance == Variance.OUT_VARIANCE) ||
        (outerVariance == Variance.OUT_VARIANCE && innerVariance == Variance.IN_VARIANCE)

/**
 * When substituting `innerVariance T` into the context `outerVariance []`, returns the variance
 * part of the result `resultVariance T`. We already know they don't conflict.
 */
private fun combineVariance(outerVariance: Variance, innerVariance: Variance) =
    when {
        outerVariance != Variance.INVARIANT -> outerVariance
        innerVariance != Variance.INVARIANT -> innerVariance
        else -> Variance.INVARIANT
    }

private fun subProjectedType(
    substitutionMap: Map<IrTypeParameterSymbol, IrTypeArgument>,
    t: IrSimpleType,
    outerVariance: Variance
): IrTypeArgument =
    substitutionMap[t.classifier]?.let { substitutedTypeArg ->
        if (substitutedTypeArg is IrTypeProjection) {
            if (conflictingVariance(outerVariance, substitutedTypeArg.variance))
                IrStarProjectionImpl
            else {
                val newProjectedType =
                    substitutedTypeArg.type.let {
                        if (t.isNullable()) it.codeQlWithHasQuestionMark(true) else it
                    }
                val newVariance = combineVariance(outerVariance, substitutedTypeArg.variance)
                makeTypeProjection(newProjectedType, newVariance)
            }
        } else {
            substitutedTypeArg
        }
    } ?: makeTypeProjection(t.substituteTypeArguments(substitutionMap), outerVariance)

private fun IrTypeArgument.upperBound(context: IrPluginContext) =
    when (this) {
        is IrStarProjection -> context.irBuiltIns.anyNType
        is IrTypeProjection ->
            when (this.variance) {
                Variance.INVARIANT -> this.type
                Variance.IN_VARIANCE ->
                    if (this.type.isNullable()) context.irBuiltIns.anyNType
                    else context.irBuiltIns.anyType
                Variance.OUT_VARIANCE -> this.type
            }
        else -> context.irBuiltIns.anyNType
    }

private fun IrTypeArgument.lowerBound(context: IrPluginContext) =
    when (this) {
        is IrStarProjection -> context.irBuiltIns.nothingType
        is IrTypeProjection ->
            when (this.variance) {
                Variance.INVARIANT -> this.type
                Variance.IN_VARIANCE -> this.type
                Variance.OUT_VARIANCE ->
                    if (this.type.isNullable()) context.irBuiltIns.nothingNType
                    else context.irBuiltIns.nothingType
            }
        else -> context.irBuiltIns.nothingType
    }

fun IrType.substituteTypeAndArguments(
    substitutionMap: Map<IrTypeParameterSymbol, IrTypeArgument>?,
    useContext: KotlinUsesExtractor.TypeContext,
    pluginContext: IrPluginContext
): IrType =
    substitutionMap?.let { substMap ->
        if (this is IrSimpleType) {
            val typeClassifier = this.classifier
            substMap[typeClassifier]?.let {
                when (useContext) {
                    KotlinUsesExtractor.TypeContext.RETURN -> it.upperBound(pluginContext)
                    else -> it.lowerBound(pluginContext)
                }
            } ?: this.substituteTypeArguments(substMap)
        } else {
            this
        }
    } ?: this

object RawTypeAnnotation {
    // Much of this is taken from JvmGeneratorExtensionsImpl.kt, which is not easily accessible in
    // plugin context.
    // The constants "kotlin.internal.ir" and "RawType" could be referred to symbolically, but they
    // move package
    // between different versions of the Kotlin compiler.
    val annotationConstructor: IrConstructorCall by lazy {
        val irInternalPackage = FqName("kotlin.internal.ir")
        val parent =
            IrExternalPackageFragmentImpl(
                DescriptorlessExternalPackageFragmentSymbol(),
                irInternalPackage
            )
        val annoClass =
            IrFactoryImpl.buildClass {
                    kind = ClassKind.ANNOTATION_CLASS
                    name = irInternalPackage.child(Name.identifier("RawType")).shortName()
                }
                .apply {
                    createImplicitParameterDeclarationWithWrappedDescriptor()
                    this.parent = parent
                    addConstructor { isPrimary = true }
                }
        val constructor = annoClass.constructors.single()
        IrConstructorCallImpl.fromSymbolOwner(constructor.constructedClassType, constructor.symbol)
    }
}

fun IrType.toRawType(): IrType =
    when (this) {
        is IrSimpleType -> {
            when (val owner = this.classifier.owner) {
                is IrClass -> {
                    if (this.arguments.isNotEmpty())
                        this.addAnnotations(listOf(RawTypeAnnotation.annotationConstructor))
                    else this
                }
                is IrTypeParameter -> owner.superTypes[0].toRawType()
                else -> this
            }
        }
        else -> this
    }

fun IrClass.toRawType(): IrType {
    val result = this.typeWith(listOf())
    return if (this.typeParameters.isNotEmpty())
        result.addAnnotations(listOf(RawTypeAnnotation.annotationConstructor))
    else result
}

fun IrTypeArgument.withQuestionMark(b: Boolean): IrTypeArgument =
    when (this) {
        is IrStarProjection -> this
        is IrTypeProjection ->
            this.type.let {
                when (it) {
                    is IrSimpleType ->
                        if (it.isNullable() == b) this
                        else makeTypeProjection(it.codeQlWithHasQuestionMark(b), this.variance)
                    else -> this
                }
            }
        else -> this
    }

typealias TypeSubstitution = (IrType, KotlinUsesExtractor.TypeContext, IrPluginContext) -> IrType

private fun matchingTypeParameters(l: IrTypeParameter?, r: IrTypeParameter): Boolean {
    if (l === r) return true
    if (l == null) return false
    // Special case: match List's E and MutableList's E, for example, because in the JVM lowering
    // they will map to the same thing.
    val lParent = l.parent as? IrClass ?: return false
    val rParent = r.parent as? IrClass ?: return false
    val lJavaId = getJavaEquivalentClassId(lParent) ?: lParent.classId
    return (getJavaEquivalentClassId(rParent) ?: rParent.classId) == lJavaId && l.name == r.name
}

// Returns true if type is C<T1, T2, ...> where C is declared `class C<T1, T2, ...> { ... }`
fun isUnspecialised(
    paramsContainer: IrTypeParametersContainer,
    args: List<IrTypeArgument>,
    logger: Logger
): Boolean {
    return isUnspecialised(paramsContainer, args, logger, paramsContainer)
}

private fun isUnspecialised(
    paramsContainer: IrTypeParametersContainer,
    args: List<IrTypeArgument>,
    logger: Logger,
    origParamsContainer: IrTypeParametersContainer
): Boolean {
    val unspecialisedHere =
        paramsContainer.typeParameters.zip(args).all { paramAndArg ->
            (paramAndArg.second as? IrTypeProjection)?.let {
                // Type arg refers to the class' own type parameter?
                it.variance == Variance.INVARIANT &&
                    matchingTypeParameters(
                        it.type.classifierOrNull?.owner as? IrTypeParameter,
                        paramAndArg.first
                    )
            } ?: false
        }
    val remainingArgs = args.drop(paramsContainer.typeParameters.size)

    val parentTypeContainer =
        paramsContainer.parents.firstIsInstanceOrNull<IrTypeParametersContainer>()

    val parentUnspecialised =
        when {
            remainingArgs.isEmpty() -> true
            parentTypeContainer == null -> {
                logger.error(
                    "Found more type arguments than parameters: ${origParamsContainer.kotlinFqName.asString()}"
                )
                false
            }
            else -> isUnspecialised(parentTypeContainer, remainingArgs, logger, origParamsContainer)
        }
    return unspecialisedHere && parentUnspecialised
}

// Returns true if type is C<T1, T2, ...> where C is declared `class C<T1, T2, ...> { ... }`
fun isUnspecialised(type: IrSimpleType, logger: Logger) =
    (type.classifier.owner as? IrClass)?.let { isUnspecialised(it, type.arguments, logger) }
        ?: false
